# -*- coding: utf-8 -*-
"""
@Time    : 2024/8/27 11:30 
@Author  : ZhangShenao 
@File    : 2.使用LangChain实现多模态.py
@Desc    : 使用LangChain实现多模态

实现步骤:
1. 调用多模态API,上传图片,让LLM推理出图片所在的城市
2. LLM获取城市信息,推理出需要调用工具GaodeWeatherTool,生成调用参数信息
3. 程序获取需要调用的工具和参数信息,实际调用工具,获取所在城市的天气情况
4. 将工具返回的天气情况作为Prompt,传递给LLM,由LLM整理成易于人类理解的格式,最终输出
"""
import os

import dotenv
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnablePassthrough
from langchain_openai import ChatOpenAI

from gaode_weather_tool import GaodeWeatherTool

# 加载环境变量
dotenv.load_dotenv()

# 创建工具列表
weather_tool = GaodeWeatherTool()

# 创建LLM,并绑定工具列表
llm = ChatOpenAI(model="gpt-4o", temperature=0, openai_api_base=os.getenv("OPENAI_API_BASE"))
llm_with_tools = llm.bind_tools(tools=[weather_tool], tool_choice="auto")

# 构造天气Prompt
weather_prompt = ChatPromptTemplate.from_messages([
    # 通过传入列表,构造多模态Prompt
    ("human", [
        {"type": "text", "text": "请根据上传的图片，获取对应城市的天气信息"},
        {"type": "image_url", "image_url": {"url": "{image_url}"}}
    ])
])

# 构造最终Prompt
final_prompt = ChatPromptTemplate.from_template("""请整理下输入的城市天气信息，并以用户友好的形式输出。
以下是原始天气信息：
<weather>
{weather}
</weather>
""")

# 构造Chain
chain = ({"weather": ({"image_url": RunnablePassthrough()} |
                      weather_prompt |
                      llm_with_tools |
                      (lambda resp: resp.tool_calls[0].get("args")) |
                      weather_tool)}
         | final_prompt | llm | StrOutputParser())

# 执行Chain,获取结果
result = chain.invoke("https://img1.baidu.com/it/u=644490943,1781886584&fm=253&fmt=auto&app=138&f=JPEG")
print(result)
